'''
Created on Oct 20, 2012

@author: Andre Dozier
'''

import glob
import string 
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm 
import access 
import sys
import scipy.stats as stats
import math 

def hours(s):
    '''Converts a string in the output files to hours'''
    ind = s.find(":")
    if ind < 0: 
        return 0.0
    return float(s[:ind]) / 24.0 + float(s[ind + 1:]) / 24.0 / 60.0

def copy(listOfVals): 
    if isinstance(listOfVals, list): 
        return [l * 1 for l in listOfVals]
    else: 
        return l * 1
    
def AF_to_kcfs(listOfVals): 
    factor = 43560.0 / 3600000.0
    if isinstance(listOfVals, list): 
        return convertlist(listOfVals,factor)
    else: 
        listOfVals[:] = listOfVals * factor
        return listOfVals * factor
    
def convertlist(listOfVals, factor):
    listOfVals[:] = [listOfVals[i] * factor for i in range(len(listOfVals))] 
    return listOfVals

def listOfArraysTo2DArray(data): 
    if not isinstance(data, list): 
        raise Exception('data needs to be a list')
    if len(data) == 0: 
        return None
    alldata = data[0]
    for d in data[1:]: 
        alldata = np.hstack((alldata, d))
    return np.array(alldata, dtype=np.float64)
    
def getLinkOutput(mdb, name): 
    '''
    Returns link flows
    '''
    if isinstance(name, list): 
        name = string.join(name, "' or LinksInfo.LName='") 
    data = np.array(mdb.execute("select Sum(Flow) from LinksOutput inner join LinksInfo ON LinksOutput.LNumber=LinksInfo.LNumber where LinksInfo.LName='%s' group by LinksOutput.TSIndex order by LinksOutput.TSIndex" % name))
    return data

def getHydroTargets(mdb, name, attributes):
    '''
    Returns the energy targets, production, and difference 
    in columns 0, 1, and 2, respectively
    
    @param attributes: 
    Can be any combination of the following (separated by a comma)
        - EnergyTarget, EnergyProduction, EnergyDifference 
    '''
    data = np.array(mdb.execute("select %s from HydroTargetOutput inner join HydroTargetsInfo ON HydroTargetOutput.HydroTargetID=HydroTargetsInfo.HydroTargetID where HydroTargetsInfo.HydroTargetName='%s' order by HydroTargetOutput.TSIndex" % (attributes, name)))
    return data 

def getHydroUnits(mdb, name, attributes):
    '''
    Returns hydro unit discharge, head, efficiency, power, and energy
    production in columns 0-4, respectively
    
    @param attributes: 
    Can be any combination of the following (separated by a comma)
        - Discharge, Head, Efficiency, Power, Energy  
    '''
    data = np.array(mdb.execute("select %s from HydroUnitOutput inner join HydroUnitsInfo ON HydroUnitOutput.HydroUnitID=HydroUnitsInfo.HydroUnitID where HydroUnitsInfo.HydroUnitName='%s' order by HydroUnitOutput.TSIndex" % (attributes, name)))
    return data

def getStorage(mdb, name, targets=False, includeBeg=True):
    if targets:
        storage = np.array(mdb.execute("select Stor_Beg,Stor_Trg from RES_STOROutput inner join NodesInfo ON RES_STOROutput.NNo=NodesInfo.NNumber where NodesInfo.NName='%s' order by RES_STOROutput.TSIndex" % name))
    else: 
        storage = np.array(mdb.execute("select Stor_Beg,Stor_End from RES_STOROutput inner join NodesInfo ON RES_STOROutput.NNo=NodesInfo.NNumber where NodesInfo.NName='%s' order by RES_STOROutput.TSIndex" % name))
    tmp = storage[0:1, 0:1]
    storage = storage[:, 1:2]
    if includeBeg: 
        storage = np.vstack((tmp, storage))
    return storage

def getResOutput(mdb, name, attributes):
    '''
    Returns a combination of attributes from the RESOutput table in the
    output database.
    
    @param attributes: 
    A comma separated string of the following columns in the table
        - Spills, Evap_Loss, Seepage, Ups_Rel, Pump_In, Gw_In, Dws_Rel,
          Pump_Out, Hydra_Cap, Head_Avg, Powr_Avg, Energy, Pwr_2nd,
          Hydro_State
    '''
    data = np.array(mdb.execute("select %s from RESOutput inner join NodesInfo ON RESOutput.NNo=NodesInfo.NNumber where NodesInfo.NName='%s' order by RESOutput.TSIndex" % (attributes, name)))
    return data    
    
def getAllData(pattern, origtargfile):
    # sorted files names 
    fnames = glob.glob(pattern)
    if fnames == []:
        raise Exception('No files found using pattern %s' % pattern)
    fnames.sort()
    
    # get originals 
    originals = np.recfromcsv(origtargfile)

    # initialize lists
    allpowerflows = []; 
    allSwings = []; allSwingTargs = [] 
    allnames = []; hydroWeights = []
    maxPFs = []; windPs = []; windTSs = []
    gclstor = []; chjstor = []; bnkstor = [] # storage
    gcldchrg = []; chjdchrg = []; bnkpumpdchrg = []; bnkgendchrg = [] # discharge
    gclGenTarg = []; chjGenTarg = []; bnkPumpTarg = []; bnkGenTarg = [] # targets and production
    gclInf = []; gclOut = []; gclSpill = []; gclDwnOut = []
    powerflowTargs = None

    # list file names 
    for fname in fnames:
    
        # read the mdb files
        splitPath = os.path.split(fname) 
        runName = splitPath[1].replace("_power.csv", "")
        mdbPath = os.path.join(splitPath[0], runName + "OUTPUT.mdb")
        if os.path.exists(mdbPath):
            mdb = access.db(mdbPath)
            gclstor.append(getStorage(mdb, "GCL"))
            chjstor.append(getStorage(mdb, "CHJ"))
            bnkstor.append(getStorage(mdb, "BNK"))
            gcldchrg.append(getHydroUnits(mdb, "GCL_PowerPlant", "Discharge"))
            chjdchrg.append(getHydroUnits(mdb, "CHJ_hydroUnit", "Discharge"))
            bnkpumpdchrg.append(getHydroUnits(mdb, "BNK_PumpUnit", "Discharge"))
            bnkgendchrg.append(getHydroUnits(mdb, "BNK_GenUnit", "Discharge"))
            gclGenTarg.append(getHydroTargets(mdb, "GCL_GenTarget", "EnergyTarget, EnergyProduction"))
            chjGenTarg.append(getHydroTargets(mdb, "CHJ_Target", "EnergyTarget, EnergyProduction"))
            bnkPumpTarg.append(getHydroTargets(mdb, "BNK_PumpTarget", "EnergyTarget, EnergyProduction"))
            bnkGenTarg.append(getHydroTargets(mdb, "BNK_GenTarget", "EnergyTarget, EnergyProduction"))

            # water flows in and out of Grand Coulee
            spill = getResOutput(mdb, "GCL", "Spills")
            gclInf.append(getLinkOutput(mdb, ["GCL_IncInf_GCL","GCL_Inflow_GCL","BNK_GCL"]) - getLinkOutput(mdb, "GCL_GCL_IncLoss"))
            gclOut.append(spill + getLinkOutput(mdb, ["GCL_Spill","GCL_PwrRel","GCL_BNK"]))
            gclDwnOut.append(spill + getLinkOutput(mdb, ["GCL_Spill","GCL_PwrRel"]))
            gclSpill.append(spill + getLinkOutput(mdb, "GCL_Spill"))
            mdb.close()
            
        # weights, transmission capacities, and wind power scaling factors
        print "Working on " + runName
        hydroI = string.find(runName, "_H_")
        maxPFI = string.find(runName, "_M_")
        windPI = string.find(runName, "_S_") 
        lastI = len(runName) + 1
        hydroWeight = float(runName[hydroI + 3:maxPFI])
        maxPF = int(runName[maxPFI + 3:windPI])
        if maxPF <= 0: 
            maxPF = float("inf")
        windP = float(runName[windPI + 3:lastI])
        
        # read file data 
        f = open(fname, "r") 
        header = f.readline() 
        names = header.strip().split(',')
        data = np.loadtxt(f, delimiter=',')
        f.close()
        
        # load the text file into an array    
        nsteps = len(data)
        swingTargs = data[:, 11]
        if powerflowTargs == None: 
            powerflowTargs = np.zeros(data.shape[0])
        powerflows = np.vstack((data[:,2] - originals.chjgen[:len(data)], data[:, 1] - originals.chjgen[:len(data)])).T
        swings = data[:, 14] - (data[:, 1] + data[:, 3] + data[:, 5] + data[:, 7])
        
        # store in list 
        windTSs.append(data[:, 9])
        allpowerflows.append(powerflows)
        allSwingTargs.append(swingTargs)
        allSwings.append(swings)
        allnames.append(runName)
        hydroWeights.append(hydroWeight) 
        maxPFs.append(maxPF)
        windPs.append(windP)
    
    return nsteps, originals, windTSs, allpowerflows, powerflowTargs, allSwingTargs, allSwings, allnames, hydroWeights, maxPFs, windPs, \
           gclstor,chjstor,bnkstor,gcldchrg,chjdchrg,bnkpumpdchrg,bnkgendchrg,gclGenTarg,chjGenTarg,bnkPumpTarg,bnkGenTarg, \
           gclInf, gclOut, gclDwnOut, gclSpill

def plot(name, data, hist=None, histI=None, bounds=None, verbose=False):
    plt.figure()
    plt.clf()
    lines = [] 
    done = []  
    names = [] 
    doneWeights = []
    doneScalar = []
    line_shapes = ['-','--',':',':-']
    for i in range(len(data)):
        if tptrn == "": 
            color = {1500: 'k', 1000: 'r', 750: 'b', 500: 'g'}.get(maxPFs[i], 'k')
            extra = maxPFs[i] in done
        elif sptrn == "":
            color = '' #auto color {1e+8: 'b', 1e+6: 'r', 1e+5: 'g'}.get(hydroWeights[i], 'b')
            extra = windPs[i] in doneScalar
        elif hptrn == "": 
            color = {1e+8: 'b', 1e+7: 'y', 1e+6: 'r', 1e+5: 'g', 1e+4: 'c', 1e+3: 'm'}.get(hydroWeights[i], 'b')
            extra = hydroWeights[i] in doneWeights
        else: 
            color = {2500: 'b', 2000: 'y', 1500: 'r', 1000: 'g', 750: 'c', 500: 'm'}.get(maxPFs[i], 'k')
            extra = maxPFs[i] in done
        
        if len(data[i].shape) >= 2:
            for j in range(data[i].shape[1]): 
                l = plt.plot(data[i][:,j], color + line_shapes[j])[0]
        else:
            l = plt.plot(data[i], color)[0]
        if not extra: 
            if tptrn == "": 
                names.append(str(maxPFs[i]) + ' MW Capacity')
                lines.append(l)
            elif sptrn == "": 
                names.append("Scalar = " + str(windPs[i]))
                lines.append(l)
            elif hptrn == "": 
                names.append(str(hydroWeights[i]) + ' hydro weight')
                lines.append(l)
            else:  
                names.append(str(maxPFs[i]) + ' MW Capacity')
                lines.append(l)
            done.append(maxPFs[i])
            doneWeights.append(hydroWeights[i])
            doneScalar.append(windPs[i])

#        if tptrn == "":
#            lines.pop(3); lines.pop(0); 
#            names.pop(3); names.pop(0); 
#            l = lines.pop(3);   n = names.pop(3)
#            lines.insert(1, l); names.insert(1, n)
#        elif sptrn == "": 
#            names.append("Scalar = " + str(windPs[i]))
#            lines.append(l)
#        elif hptrn == "": 
#            names.append(str(hydroWeights[i]) + ' hydro weight')
#            lines.append(l)
    
    if hist != None: 
        names.append("Historical")
        if histI != None:
            lines.append(plt.plot(histI[:nsteps], hist[:nsteps], 'x', alpha=0.5, linewidth=2)[0])
        else:
            lines.append(plt.plot(hist[:nsteps], 'k', alpha=0.5, linewidth=2)[0])
            
    if bounds != None:
        plt.gca().set_ybound(bounds)
         
    plt.ylabel(name, fontsize=14) 
    plt.xlabel("TimeStep", fontsize=14)
    prop = fm.FontProperties(size=12)
    plt.legend(lines, names, prop=prop, loc='best')
    basepath = os.path.join(outdir, name)
    plt.savefig(basepath + theext, dpi=300)
    alldata = listOfArraysTo2DArray(data)
    if verbose: 
        indices = np.argwhere(alldata == None)
        print alldata
        print indices 
        print alldata[indices]
    if alldata != None: 
        np.savetxt(basepath + "_data.txt", alldata, delimiter='\t')
    if hist != None: 
        np.savetxt(basepath + "_hist.txt", hist, delimiter='\t')

def total_dissolved_gases():
    # total dissolved gases
    import modelpy as mpy
    import datapy as d 
    
    # get the model 
    f = open(r"gclwqmodel.csv", "r")
    header = f.readline() 
    names = header.strip().split(',')
    linMinfo = np.loadtxt(f, usecols=[1, 2, 3], delimiter=",")
    f.close()
    w = linMinfo[:, 2:3]
    Xmeans = linMinfo[1:, 1:2].T
    Xstd = linMinfo[1:, 0:1].T
    Ymean = float(names[1])
    Ystd = float(names[2])
    Xstdizer = d.Standardizer(means=Xmeans, stdevs=Xstd)
    Ystdizer = d.Standardizer(means=Ymean, stdevs=Ystd)
    tdgmodel = mpy.LinearModel(w, Xstdizer, Ystdizer)

    # get the data 
    f = open(r"gclwqdata.csv", "r")
    header = f.readline() 
    names = header.strip().split(',')
    Xtrain = np.loadtxt(f, delimiter=",")
    tdgtrg = Xtrain[:, -1:]
    tdgIs = np.arange(tdgtrg.shape[0]).reshape(-1, 1)
    tdgGT0 = (tdgtrg > 0.0)
    tdgtrg = tdgtrg[tdgGT0]
    tdgtrg = tdgtrg.reshape(-1, 1)
    tdgIs = tdgIs[tdgGT0]
    tdgIs = tdgIs.reshape(-1, 1)

    Xtrain = Xtrain[:, 2:-1]
    
    # model with new data
    def gclElev(Volume): 
        return -0.00000000000083851345 * Volume ** 2.0 + \
            0.00001991247749129520 * Volume + \
            1208.73167352232000000000
    def gclTWelev(dwnstreamout):
        return 0.00000000000000000296 * dwnstreamout ** 4.0 - \
            0.00000000000049175166 * dwnstreamout ** 3.0 + \
            0.00000001951356586505 * dwnstreamout ** 2.0 + \
            0.00086276849624766800 * dwnstreamout + \
            952.10464891132900000000
    
    # get the other data 
    def avgElev(beg, end): 
        return (gclElev(beg) + gclElev(end)) / 2.0 
    fb_elev = [avgElev(stor[:-1], stor[1:]) for stor in gclstor] 
    tw_elev = [gclTWelev(outflow) for outflow in gclDwnOut] 
    gclSpillSqr = [spill ** 2 for spill in gclSpill]
    
    # get all the Xtrain matrices for simulating TDGs for each scenario
    cols = [0, 1, 8, 9, 10, 11]
    def newX(newColumns):
        tmp = 1.0 * Xtrain[:newColumns.shape[0]]
        tmp[:, cols] = newColumns 
        return tmp[:]
    allX = [newX(np.hstack((fb_elev[i], tw_elev[i],
                            gclInf[i], gclOut[i],
                            gclSpill[i], gclSpillSqr[i]))) 
            for i in range(len(gclSpill))]
    
    # simulate and plot TDG 
    tdgs = [tdgmodel.Use(allX[i]) for i in range(len(allX))]
    plot("Total dissolved gas (%)", tdgs, tdgtrg, tdgIs)
    return tdgs 
    
def plot_against_wind(allSwings, ylabel, shortname, figsize=None,xlims=None,ylims=None,legend_loc='best'): 

    # start figure 
    plt.figure(figsize=figsize)
    
    
    # build data 
    names = []
    lines = []
    pfs = list(set(maxPFs))
    pfs.sort()
    
    # summary data
    allimports = []
    allwind = []
    allwind_std = []
    allwind_n = []
    
    # colors 
    allcolors = ['b','g','r','c','m','y','brown','blueviolet','orange',
                 'palegreen','olive','wheat','salmon','lightslategray']
    colors = [allcolors[i % len(allcolors)] for i in range(len(pfs))] 
    for currPF, color in zip(pfs, colors): 
        # imports 
        imports = np.array([np.abs(allSwings[i]).mean() for i in range(len(allSwings)) if maxPFs[i] == currPF])
        
        # wind 
        wind = np.array([windTSs[i].mean() for i in range(len(windTSs)) if maxPFs[i] == currPF])
        wind_n = np.array([len(windTSs[i]) for i in range(len(windTSs)) if maxPFs[i] == currPF])
        wind_std = np.array([windTSs[i].std() for i in range(len(windTSs)) if maxPFs[i] == currPF])
        wind_std = wind_std / np.sqrt(wind_n)

        # stats on imports
        # cibounds = np.asarray([stats.bayes_mvs(np.abs(allSwings[i]), alpha=0.95)[0][1] for i in range(len(allSwings)) if maxPFs[i] == currPF])
        # imports2_5 = cibounds[:,0]
        # imports97_5 = cibounds[:,1]
        imports_std = np.array([np.abs(allSwings[i]).std()/math.sqrt(len(allSwings[i])) for i in range(len(allSwings)) if maxPFs[i] == currPF])
        imports2_5 = np.asarray([stats.norm.ppf(0.025,loc=mu, scale=sig) for mu,sig in zip(imports,imports_std)])
        imports97_5 = np.asarray([stats.norm.ppf(0.975,loc=mu, scale=sig) for mu,sig in zip(imports,imports_std)])
        
        # sort values by wind factor 
        ind = np.argsort(wind)
        wind = wind[ind]
        wind_std = wind_std[ind]
        wind_n = wind_n[ind]
        imports = imports[ind]
        imports2_5 = imports2_5[ind]
        imports97_5 = imports97_5[ind]
        
        # plot and name
        lines.extend(plt.plot(wind, imports, color=color))
        plt.fill_between(wind, imports2_5, imports97_5, facecolor=color, alpha=0.5)
        names.append(str(currPF) + " MW Capacity")
        
        # add to summary lists
        allimports.append(imports2_5)
        allwind.append(wind)
        allwind_std.append(wind_std)
        allwind_n.append(wind_n)

    if xlims != None: 
        plt.xlim(xlims)
    else:
        xlims = plt.xlim()
    if ylims != None: 
        plt.ylim(ylims)
    else: 
        ylims = plt.ylim()
        
    # analyze where the averages cross the biggest transmission scenario 
    if len(pfs) > 1: 
        # a function to calculate potential wind power that can be integrated into the system
        def calcWindPotential(wind, imports, threshold):
            right_i = np.argwhere(imports > threshold)
            if len(right_i) == 0: 
                return None, None
            right_i = int(right_i[0])
            left_i = right_i - 1
            right_x = wind[right_i]
            left_x = wind[left_i]
            right_y = imports[right_i]
            left_y = imports[left_i]
            Pwind = (right_x-left_x)/(right_y-left_y) * (threshold-left_y) + left_x
            if abs(Pwind - left_x) < abs(Pwind - right_x): 
                wind_index = left_i
            else: 
                wind_index = right_i
            return wind_index, Pwind
            
        # loop through all but the last one
        threshold = imports97_5[:6].mean() # threshold is the upper bound of the CI of the "no transcap scenario"
        lines.extend(plt.plot(xlims, [threshold]*2, 'k--'))
        names.append('Threshold')
        Pwind_means = []
        Pwind_stds = []
        Pwind_n = []
        for wind, imports, wind_std, wind_n, color in zip(allwind, allimports, allwind_std, allwind_n, colors): 
            i, Pwind = calcWindPotential(wind, imports, threshold)
            if i == None: 
                break
            wind2_5 = stats.norm.ppf(0.025, loc=Pwind, scale=wind_std[i])
            wind97_5 = stats.norm.ppf(0.975, loc=Pwind, scale=wind_std[i])
            Pwind_means.append(Pwind)
            Pwind_stds.append(wind_std[i])
            Pwind_n.append(wind_n[i])
            yloc = 0.005 * (ylims[1]-ylims[0]) + ylims[0]
            ywidth = 0.01 * (ylims[1]-ylims[0]) + ylims[0]
            plt.plot(Pwind, yloc, color=color, linestyle='', marker='|')
            x = [wind2_5, wind97_5]
            y1 = [yloc-ywidth/2.]*2
            y2 = [yloc+ywidth/2.]*2
            plt.fill_between(x,y1,y2,facecolor=color,edgecolor='',alpha=0.5)
        
        # compare each potential wind power against each other
        print 'Pwind Mean (95% CI):'
        for i in range(len(Pwind_means)-1): 
            # left side model
            currpf = pfs[i]
            currPwind_mean = Pwind_means[i]
            currPwind_std = Pwind_stds[i]
            currPwind_n = Pwind_n[i]
            
            # right side model 
            nextpf = pfs[i+1]
            nextPwind_mean = Pwind_means[i+1]
            nextPwind_std = Pwind_stds[i+1]
            nextPwind_n = Pwind_n[i+1]
            
            # model of difference  
            Pmean = nextPwind_mean-currPwind_mean
            Pstd = math.sqrt(currPwind_std ** 2 / currPwind_n + nextPwind_std ** 2 / nextPwind_n)
            n = stats.norm(loc=Pmean, scale=Pstd)
            CI = (n.ppf(0.025),n.ppf(0.975))
            contains = 'contains zero' if np.product(CI) < 0 else 'does not contain zero' 
            print '  %g MW scen minus %g MW scen:' % (nextpf, currpf)
            print '     %.0f (%.0f, %.0f) - %s' % (Pmean, CI[0], CI[1], contains)

    # plot values in the figure
    plt.xlabel("Average wind power penetration (MW)", fontsize=14)
    plt.ylabel(ylabel, fontsize=14)
    print 'xlims', xlims
    print 'ylims', ylims
    plt.xlim(xlims)
    plt.ylim(ylims)
    plt.tick_params(axis='both',which='major',labelsize=10)
    plt.legend(lines, names, loc=legend_loc, fontsize=10)
    plt.savefig(os.path.join(outdir, "Wind vs " + shortname + theext), dpi=300)


#=======================================================================
# Main method
#=======================================================================
if __name__ == "__main__":

    # plot types 
    plottypes = ['tstep', 'tdg', 'imports', 'storage', 'all']
    def plotstr(s): 
        s = s.lower() 
        if not s in plottypes: 
            raise Exception('Flag --plot %s not recognized! Must be a string within the following list: %s' % (s,', '.join(plottypes)))
        return s
            
    # arguments
    import argparse
    parser = argparse.ArgumentParser(description="Process MODSIM outputs")
    parser.add_argument('scen', type=str, help="Scenario name (and output subfolder)")
    parser.add_argument('--tcap', type=int, default=None, help="Fix the transmission capacity to analyze")
    parser.add_argument('--weight', type=int, default=None, help="Fix the weight to analyze")
    parser.add_argument('--wind', type=int, default=None, help="Fix the wind scaling factor to analyze")
    parser.add_argument('--ext', type=str, default='.png', help="File extension of output figures")
    parser.add_argument('--plot', type=plotstr, nargs='+', default='all', help="Plot type must be in the following list ('all' is default): {0}".format(', '.join(plottypes)))
    parser.add_argument('--xlims', type=float, nargs=2, default=None, help="Lower and upper limit on the x axis for the import-wind plot")
    parser.add_argument('--ylims', type=float, nargs=2, default=None, help="Lower and upper limit on the y axis for the import-wind plot")
    parser.add_argument('--figsize', type=float, nargs=2, default=None, help="Width and Heighth of the figure")
    parser.add_argument('--legend', type=str, default='best', help="Location of the legend. Default is 'best'")
    parser.add_argument('-d','--dirname',type=str,default=None, help="Directory containing the MODSIM output files to analyze")
    args = parser.parse_args()

    # close all plots 
    plt.close('all')
    
    def eformat(val, exp_digits):
        s = "%.0e"%(val)
        mantissa, exp = s.split('e')
        # add 1 to digits as 1 is taken by sign +/-
        return "%se%+0*d" % (mantissa, exp_digits+1, int(exp))
    
    # get inputs from arguments 
    scenName = args.scen
    transCap = args.tcap
    hweight = args.weight
    sptrn = hptrn = tptrn = "" 
    if transCap != None: 
        tptrn = "*_M_%d" % transCap
    if hweight != None: 
        hptrn = "*_H_" + eformat(hweight, 1) 
    if args.wind != None: 
        sptrn = "*_S_%d" % args.wind
    theext = args.ext
    do_tsplots = 'all' in args.plot or 'tstep' in args.plot
    do_tdgplots = 'all' in args.plot or 'tdg' in args.plot
    do_impplots = 'all' in args.plot or 'imports' in args.plot
    do_storplots = 'all' in args.plot or 'storage' in args.plot
    xlims = args.xlims
    ylims = args.ylims 
    figsize = args.figsize
    legend_loc = args.legend
    dirname = args.dirname
    
    # directories
    if dirname != None: 
        mydir = dirname
    else: 
        # list directories
        mydir = os.path.join('..','GCL_TwoBus','*')
        thedirs = glob.glob(mydir)
        thedirs = [d for d in thedirs if os.path.isdir(d)]
        thedirs.sort()
        ids = range(len(thedirs))
        
        # query for directory
        for (i,d) in zip(ids,thedirs):
            print '%5d:  %s' % (i,d)
        mydir_i = raw_input('\nWhat directory would you like to work in?\n' + 
                              '    Select a number from the list above (q to quit): ')
        if mydir_i.lower() == 'q': 
            sys.exit()

        # set the directory name
        try:
            mydir_i = int(mydir_i)
            mydir = thedirs[mydir_i]
        except:
            raise Exception('Must be a number from %d to %d.' % (ids[0], ids[-1])) 
                
    print '\nWorking in directory: %s\n' % os.path.relpath(mydir)
    
    # directories
    origtargfile = r"original_targets.csv"
    mypattern = os.path.join(mydir, hptrn + tptrn + sptrn + r"*_power.csv")
    outdir = os.path.join(mydir, scenName)
    if not os.path.exists(outdir): 
        os.makedirs(outdir)

    # get all the data 
    nsteps, originals, windTSs, allpowerflows, powerflowTargs, allSwingTargs, allSwings, allnames, hydroWeights, maxPFs, windPs, \
        gclstor,chjstor,bnkstor,gcldchrg,chjdchrg,bnkpumpdchrg,bnkgendchrg,gclGenTarg,chjGenTarg,bnkPumpTarg,bnkGenTarg, \
        gclInf, gclOut, gclDwnOut, gclSpill \
        = getAllData(mypattern, origtargfile)
        
    if do_tsplots: 
        # power flows 
        print 'Plotting power flows...'
        plot("Power flow (MW)", allpowerflows, powerflowTargs)
        plot("Power imports (MW)", allSwings)
        plot("Power import targets (MW)", allSwingTargs, bounds=[-1000, 6000])
        
        # storages
        print 'Plotting simulated storage...'
        plot("Grand Coulee Storage (AF)", gclstor, originals.gclstor)
        plot("Chief Joseph Storage (AF)", chjstor, originals.chjstor)
        plot("Banks Lake Storage (AF)", bnkstor, originals.bnkstor)
        
        # hydro units 
        print 'Plotting discharge...'
        plot("Grand Coulee Hydro Unit Discharge (kcfs)", AF_to_kcfs(gcldchrg))
        plot("Chief Joseph Hydro Unit Discharge (kcfs)", AF_to_kcfs(chjdchrg))
        plot("Banks Lake Hydro Pump Unit Discharge (kcfs)", AF_to_kcfs(bnkpumpdchrg))
        plot("Banks Lake Hydro Gen Unit Discharge (kcfs)", AF_to_kcfs(bnkgendchrg))
        
        # hydro targets
        print 'Plotting hydropower production and targets...'
        plot("Grand Coulee Hydro Target (MW)", gclGenTarg, originals.gclgen)
        plot("Chief Joseph Hydro Target (MW)", chjGenTarg, originals.chjgen)
        plot("Banks Lake Hydro Pump Target (MW)", bnkPumpTarg, originals.bnkpump)
        plot("Banks Lake Hydro Gen Target (MW)", bnkGenTarg, originals.bnkgen)

        # inflows / outflows 
        print 'Plotting Grand Coulee mass balance...'
        gclInfAF = copy(gclInf)
        gclOutAF = copy(gclOut)
        plot("Grand Coulee Inflow (kcfs)", AF_to_kcfs(gclInf), AF_to_kcfs(originals.gclinf))
        plot("Grand Coulee Outflow (kcfs)", AF_to_kcfs(gclOut), AF_to_kcfs(originals.gclout))
        plot("Grand Coulee Spill (kcfs)", AF_to_kcfs(gclSpill), AF_to_kcfs(originals.gclspill))
        # for i, q, s in zip(gclInfAF, gclOutAF, gclstor):
            # massbal = (s[:-1] + i - q) - s[1:]
            # if (np.abs(massbal) > 0.001).any():
                # print '   Mass balance mismatch:      ', massbal
                # print '   Mean mass balance mismatch: ', massbal.mean()

    # Analyze the total dissolved gases
    if do_tdgplots: 
        print 'Plotting Total Dissolved Gases...'
        tdgs = total_dissolved_gases()
        plot_against_wind(tdgs, "Average Total Dissolved Gases (%)", "TDG", figsize=figsize,xlims=xlims,ylims=ylims,legend_loc=legend_loc)
    
    # plot avg. wind power import against wind power penetration
    if do_impplots: 
        print 'Plotting imported power against wind power penetration...'
        plot_against_wind(allSwings, "Average power mismatch (MW)", "Imports", figsize=figsize,xlims=xlims,ylims=ylims,legend_loc=legend_loc)

    if do_storplots:
        print 'Plotting storage against wind power penetration...'
        plot_against_wind(chjstor, "Average storage (acre-feet)", "Storage", figsize=figsize,xlims=xlims,ylims=ylims,legend_loc=legend_loc)



